use pyo3::prelude::*;

///# hello world功能
/// ## args: username 用户名
/// ## return : 对输入的用户问好。
/// pyo3demo.output_hw("username")
/// return "hello username\n ——from Rust PyO3 "
#[pyfunction]
fn output_hw(user:String) ->String{
    format!("hello {} \n ——from Rust PyO3 ",user)
}

//测试方法
// #[test]
// fn test_output_hw(){
//     let user = String::from("godxia");
//     assert_eq!("hello godxia\n ——from Rust PyO3 ",output_hw(user));
// }

/**
    #两数相乘
    ## args: x : float,y: float
    ## return x * y
    ```
    res = two_multi_float(3.0,4.0)
    assert(res,12)
    ``` 
*/
#[pyfunction]
fn two_multi_float(x:f64, y:f64) -> f64{
    x*y
}

/**
    #两数相除
    ## args: x : float,y: float
    ## return x / y
    ```
    res = two_multi_float(12.0,4.0)
    assert(res,3.0)
    ``` 
*/
#[pyfunction]
fn two_divi_float(x:f64, y:f64) -> f64{
    x/y
}

/**
    #整数累加
    ## args: x : i32
    ## return res= 1+2+3...+x
    ```
    res = int_sum(10)
    assert(res,55)
    ``` 
*/
#[pyfunction]
fn int_sum(x:i32)->i128{
    let mut res =0;
    for i in 1..x+1{
        res+=i;
    }
    res as i128
}
/** 
 * # 一维矩阵相乘
    ## args: x :[i32],y[i32]
    ## return res= x * y
    ```
    x = [1,2]
    y = [3,4]
    res = matrix_multi(x,y)
    assert(res,[[3, 4], [6, 8]])
    ```
 * 
*/
#[pyfunction]
fn matrix_multi(x:Vec<i32>,y:Vec<i32>)->Vec<Vec<i32>>{
    let mut res = Vec::new();
    for x1 in x.iter(){
        let mut xarr = Vec::new();
        for y1 in y.iter(){
            xarr.push(x1*y1);
        }
        res.push(xarr)
    }
    res
}

/** 
 * # 素数判定 & 最小因素
    ## args: x :i128
    ## return 第一个被找到的因数，如果等于1，则表示是素数
    ```
    assert(is_prime(5),1)
    assert(is_prime(7),1)
    assert(is_prime(10),2)
    assert(is_prime(77),7)
    ```
 * 
*/
#[pyfunction]
fn is_prime(x:i128)->i128{
    for i in 2..(x as f64).sqrt() as i128 +1{
        if x % i ==0{
            return i
        }
    }
    1
}

/** 
 * # 查找比输入参数大的最近的一组孪生素数
    ## args: x :i128
    ## return 最近的一组孪生素数
    ```
    assert(twin_primes(5),(11,13))
    assert(twin_primes(100),(101,103))
    ```
 *
*/
#[pyfunction]
fn twin_primes(val:i128)->(i128,i128){
    let mut v = val+1;
    loop {
        if is_prime(v)==1 && is_prime(v+2)==1{
            return (v,v+2);
        }
        else{
            v =v +1;
        }
    }
}

/** 
 * # 斐波拉契数列
    ## args: x :i32
    ## return 指定代数的斐波拉契数列
    ```
    assert(fibonacci(5),[1,1,2,3,5])
    assert(is_prime(10),[1,1,2,3,5,8,13,21,34,55])
    ```
 *
*/
#[pyfunction]
fn fibonacci(sev:i32)->Vec<i32>{
    let mut v: Vec<i32> = vec![1,1];
    for i in 2..sev{
        v.push(v[i as usize-2] as i32 + v[i as usize -1] as i32);
    }
    v
}

use polars_core::prelude::*;
use polars_io::prelude::*;
use std::fs::File;


fn groupby_csv(csv: String)->Result<DataFrame,PolarsError>{
    let df = CsvReader::from_path(csv)?
            .has_header(true)
            .finish();
    match df{
        Ok(d) =>d.groupby(["DLMC"])?.select(["TBMJ"]).sum(),
        Err(e) =>panic!("Error groupb")
    }    
}

#[pyfunction]
fn grouby(csv:String) -> String{
    match groupby_csv(csv){
        Ok(d) => return d.to_string(),
        Err(_) => return "error".to_string()
    }
}

#[cfg(test)]
mod test{
    use super::*;
    #[test]
    fn two_multi_float_test(){
        assert_eq!(two_multi_float(3.0,4.0),12.0)
    }
    #[test]
    fn two_divi_float_test(){
        assert_eq!(two_divi_float(12.0,3.0),4.0)
    }
    #[test]
    fn int_sum_test(){
        assert_eq!(55_i128,int_sum(10))
    }
    #[test]
    fn matrix_multi_test(){
        assert_eq!(matrix_multi(vec![1,2],vec![3,4]),vec![vec![3,4],vec![6,8]]);
    }

    #[test]
    fn is_prime_test(){
        assert_eq!(is_prime(5),1_i128);
        assert_eq!(is_prime(10),2_i128);
    }

    #[test]
    fn twin_primes_test(){
        assert_eq!(twin_primes(5),(11,13));
    }
    #[test]
    fn fibonacci_test(){
        assert_eq!(fibonacci(5),vec![1,1,2,3,5])
    }

    #[test]
    fn grouby_test(){
        let s = grouby("e:/jbnt.csv".to_string());
        println!("{:#?}",s);
        assert!(true)
    }
    
}



//Python模块封装方法
//这个方法会把上面的pyfunction封装成可以在Python代码里面调用的模块和方法
#[pymodule]
fn pyo3demo(_py: Python<'_>, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(output_hw, m)?)?;
    m.add_function(wrap_pyfunction!(two_multi_float, m)?)?;
    m.add_function(wrap_pyfunction!(two_divi_float, m)?)?;
    m.add_function(wrap_pyfunction!(int_sum, m)?)?;
    m.add_function(wrap_pyfunction!(matrix_multi, m)?)?;
    m.add_function(wrap_pyfunction!(is_prime, m)?)?;
    m.add_function(wrap_pyfunction!(twin_primes, m)?)?;
    m.add_function(wrap_pyfunction!(fibonacci, m)?)?;
    m.add_function(wrap_pyfunction!(grouby, m)?)?;
    Ok(())
}